class InternalModel:
    def __init__(self, model_function=None):
        # 如果没有外置函数则调用内置函数
        if model_function is None:
            self.model_func = self.SEIR_func
        else:
            self.model_func = model_function

    @staticmethod
    def SEIR_func(population, model_param):
        # 判断所给参数是否合法，不合法则返回原人口组成
        if model_param['contact_infected'] < 0 or model_param['contact_exposed'] < 0 \
                or model_param['alpha'] < 0 or model_param['alpha'] > 1 \
                or model_param['beta'] < 0 or model_param['beta'] > 1 \
                or model_param['lambda_infected'] < 0 or model_param['lambda_infected'] > 1 \
                or model_param['lambda_exposed'] < 0 or model_param['lambda_exposed'] > 1 \
                or model_param['mu'] < 0 or model_param['mu'] > 1 \
                or model_param['sigma'] < 0 or model_param['sigma'] > 1:
            return population
        # N 为总人数
        N = population['susceptible'] + population['exposed'] + population['infected'] + population['recovered'] + \
            population['dead']
        population_new = {'susceptible': 0, 'exposed': 0, 'infected': 0, 'recovered': 0, 'dead': 0}
        # 分别计算两种感染源的感染人数
        infected_by_infected = model_param['contact_infected'] * model_param['lambda_infected'] * \
                               population['susceptible'] * population['infected'] / N
        infected_by_exposed = model_param['contact_exposed'] * model_param['lambda_exposed'] * \
                              population['susceptible'] * population['infected'] / N
        # 预计感染的人数超过原有的易感人数，则通过比例计算满足要求的两种感染人数
        if infected_by_exposed + infected_by_infected > population['susceptible']:
            ratio = infected_by_infected / (infected_by_exposed + infected_by_infected)
            infected_by_infected = ratio * population['susceptible']
            infected_by_exposed = population['susceptible'] - infected_by_infected
            population_new['susceptible'] = model_param['mu'] * population['recovered']
        else:
            # 当前易感者=历史易感者+康复者未产生免疫数量-感染者传染数量-潜伏者传染数量
            population_new['susceptible'] = population['susceptible'] + \
                                            model_param['mu'] * population['recovered'] - \
                                            infected_by_infected - infected_by_exposed
        # 当前潜伏者=历史潜伏者+感染者传染数量+潜伏者传染数量-潜伏者发病数量，必大于0
        population_new['exposed'] = population['exposed'] + \
                                    infected_by_infected + infected_by_exposed - \
                                    model_param['sigma'] * population['exposed']
        # 分别计算治愈和死亡人数
        recovered = model_param['alpha'] * population['infected']
        dead = model_param['beta'] * population['infected']
        # 预计死亡和治愈的人数数超过原有的感染人数，则通过比例计算满足要求的两种人数
        if recovered + dead > population['infected']:
            ratio = recovered / (recovered + dead)
            recovered = ratio * population['infected']
            dead = population['infected'] - recovered
            population_new['infected'] = model_param['sigma'] * population['exposed']
        else:
            # 当前感染者=历史感染者+潜伏者发病数量-治愈数量-死亡数量
            population_new['infected'] = population['infected'] + model_param['sigma'] * population['exposed'] - \
                                         recovered - dead
        # 当前治愈者=历史治愈者+感染者治愈人数-康复者未产生免疫数量，必大于0
        population_new['recovered'] = population['recovered'] + recovered - \
                                      model_param['mu'] * population['recovered']
        # 当前死亡者=历史死亡者+感染者死亡人数
        population_new['dead'] = population['dead'] + dead
        return population_new

    def calculate(self, population, model_param):
        return self.model_func(population, model_param)

    def set_func(self, model_func):
        self.model_func = model_func
